"""
通用多标的策略基础类
提供策略开发的统一框架，支持多品种交易
"""
import copy
import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
from sklearn.linear_model import LinearRegression

class BaseStrategy:
    strategy_name = "BaseStrategy"
    def __init__(self, api, symbol_info, market_period):
        self.api = api
        self.all_info = {}
        i_symbols = []
        # 初始化每个品种的信息
        for symbol, info in symbol_info.items():
            self.all_info[symbol] = info
            i_symbol = f"KQ.i@{info['exchange']}.{symbol}"
            self.all_info[symbol]["i_symbol"] = i_symbol
            i_symbols.append(i_symbol)
            self.all_info[symbol]["m_symbol"] = f"KQ.m@{info['exchange']}.{symbol}"
            for k, v in info.items():
                self.all_info[symbol][k] = v
        # 数据订阅 - 使用指定周期K线
        self.market_period = market_period
        self.kline = self.api.get_kline_serial(i_symbols, market_period * 60, data_length=800)
        
    def calculate_signal(self, symbol,current_pos, kline) -> float:
        """"
        用来计算信号
        输入：
        symbol标的类型
        kline该标的对应k线数据
        输出：
        0-空仓，1多头，-1空头
        """
        return 0
        
    def on_bar(self):
        """主循环 - 优化版本"""
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.kline.iloc[-1], "datetime"):
                # 一次性获取公共数据，避免重复调用
                account_balance = self.api.get_account().balance
                positions = self.api.get_position()
                symbol_count = len(self.all_info)
                
                # 遍历所有品种
                for symbol in self.all_info.keys():
                    # 获取品种信息
                    m_symbol = self.all_info[symbol]["m_symbol"]
                    symbol_info = self.api.get_quote(m_symbol)
                    main_symbol = symbol_info.underlying_symbol
                    
                    # 计算仓位大小
                    symbol_value = symbol_info.volume_multiple * symbol_info.last_price
                    position_size = int(account_balance / symbol_value / symbol_count)
                    if position_size < 1:
                        position_size = 1

                    # 获取当前持仓
                    target_pos = TargetPosTask(self.api, main_symbol)
                    main_pos = positions.get(main_symbol, None)
                    current_pos = main_pos.pos if main_pos else 0
                    
                    # 平掉非主力合约的持仓
                    for sym, pos in positions.items():
                        if symbol in sym and sym != main_symbol and pos.pos != 0:
                            TargetPosTask(self.api, sym).set_target_volume(0)

                    # 获取对应品种的K线数据 - 优化版本
                    kline = self._get_symbol_kline(symbol)
                    if kline is None or len(kline) == 0:
                        continue

                    # 计算交易信号
                    signal,score = self.calculate_signal(symbol, current_pos, kline)
                    
                    # 执行交易 - signal可以是小数，表示仓位百分比
                    target_volume = int(position_size * signal)
                    target_pos.set_target_volume(target_volume)
                    
                    # # 示例2: 在另一个附图画一根比ma小4的宽度为4的紫色指标线
                    # self.kline["score"] = score
                    # self.kline["score_b"] = self.all_info[symbol]["long_threshold"]
                    # self.kline["score_s"] = self.all_info[symbol]["short_threshold"]
                    # self.kline["score_bs"] = self.all_info[symbol]["long_close_threshold"]
                    # self.kline["score_ss"] = self.all_info[symbol]["short_close_threshold"]
                    # self.kline["score.board"] = "MA4"  # 设置为另一个附图
                    # self.kline["score_b.board"] = "MA4"  # 设置为另一个附图
                    # self.kline["score_s.board"] = "MA4"  # 设置为另一个附图
                    # self.kline["score_bs.board"] = "MA4"  # 设置为另一个附图
                    # self.kline["score_ss.board"] = "MA4"  # 设置为另一个附图
                    # self.kline["score.color"] = "#9933FF" # 0xFF9933CC  # 设置为紫色, 或者 "#9933FF"
                    # self.kline["score.width"] = 2  # 设置宽度为4，默认为1
    
    def _get_symbol_kline(self, symbol):
        """
        获取指定品种的K线数据 - 优化版本
        
        Args:
            symbol: 品种代码
            
        Returns:
            DataFrame: 该品种的K线数据，如果未找到则返回None
        """
        # 查找对应的列索引
        for i, sym in enumerate(self.all_info.keys()):
            if sym == symbol:
                # 构建列名
                add_str = "" if i == 0 else str(i)
                symbol_col = "symbol" if i == 0 else f"symbol{i}"
                
                # 检查列是否存在
                if symbol_col not in self.kline.columns:
                    continue
                
                # 验证symbol匹配
                if symbol not in str(self.kline[symbol_col].iloc[-1]):
                    continue
                
                # 构建需要的列名列表
                cols = [
                    "datetime",
                    f"id{add_str}",
                    f"open{add_str}",
                    f"high{add_str}",
                    f"low{add_str}",
                    f"close{add_str}",
                    f"volume{add_str}",
                    f"open_oi{add_str}",
                    "duration"
                ]
                
                # 检查所有列是否存在
                if not all(col in self.kline.columns for col in cols):
                    continue
                
                # 使用视图而不是深拷贝，大幅提升性能
                # iloc[:-1] 排除最后一根未完成的K线
                kline_data = self.kline[cols].iloc[:-1].copy()
                
                # 重命名列
                kline_data.columns = [
                    "datetime", "id", "open", "high", "low",
                    "close", "volume", "open_oi", "duration"
                ]
                
                return kline_data
        
        return None

